/*
 * Copyright(C) 2022. Huawei Technologies Co.,Ltd. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "model.h"
#include "math.h"
#include "fstream"

using namespace std;

namespace
{
  const uint32_t YUV_BYTE_NU = 3;
  const uint32_t YUV_BYTE_DE = 2;
  const uint32_t VPC_H_ALIGN = 2;
  const uint32_t YOLOV3_RESIZE = 416;
  const int TENSORS_PER_OBJECT = 3;
  const int BIN_WIDTH_IN_DEGREES = 3;
  const double PI = acos(-1);
  const int TWO = 2;
};

Yolo::Yolo(const V2Param &v2Param)
{
  deviceId = v2Param.deviceId;
  std::string modelPath = v2Param.modelPath;
  std::string labelPath = v2Param.labelPath;
  std::string configPath = v2Param.configPath;

  // model init
  modelDptr = std::make_shared<MxBase::Model>(modelPath, deviceId);
  if (modelDptr == nullptr)
  {
    LogError << "modelDptr nullptr";
  }

  // postprocessor init
  std::map<std::string, std::string> postConfig;
  postConfig.insert(std::pair<std::string, std::string>("postProcessConfigPath", configPath));
  postConfig.insert(std::pair<std::string, std::string>("labelPath", labelPath));
  postProcessorDptr = std::make_shared<MxBase::Yolov3PostProcess>();
  if (postProcessorDptr == nullptr)
  {
    LogError << "postProcessorDptr nullptr";
  }
  postProcessorDptr->Init(postConfig);
};

APP_ERROR Yolo::Infer(MxBase::Image &inputImage, std::vector<MxBase::Tensor> &outputs)
{
  APP_ERROR ret;

  // move image to device
  MxBase::Tensor tensorImg = inputImage.ConvertToTensor();
  ret = tensorImg.ToDevice(deviceId);
  if (ret != APP_ERR_OK)
  {
    LogError << "ToDevice failed, ret=" << ret;
    return ret;
  }

  // make infer input
  std::vector<MxBase::Tensor> inputs = {tensorImg};

  // do infer
  outputs = modelDptr->Infer(inputs);

  // move result to host
  for (size_t i = 0; i < outputs.size(); i++)
  {
    outputs[i].ToHost();
  }

  return APP_ERR_OK;
};

APP_ERROR Yolo::PostProcess(std::vector<MxBase::Tensor> &yoloOutputs, std::vector<MxBase::Rect> &cropRects,
                            int originalImageWidth, int originalImageHeight)
{
  // 后处理参数1：输入tensor
  std::vector<MxBase::TensorBase> tensors;
  for (size_t i = 0; i < yoloOutputs.size(); i++)
  {
    MxBase::MemoryData memoryData(yoloOutputs[i].GetData(), yoloOutputs[i].GetByteSize());
    MxBase::TensorBase tensorBase(memoryData, true, yoloOutputs[i].GetShape(), MxBase::TENSOR_DTYPE_INT32);
    tensors.push_back(tensorBase);
  }

  // 后处理参数2：ObjectInfo，目标检测类任务的目标框信息
  std::vector<std::vector<MxBase::ObjectInfo>> objectInfos;

  // 后处理参数3：ResizedImageInfo，用于记录图像类任务中模型前处理中的缩放方式，提供给模型后处理的坐标还原使用
  auto shape = yoloOutputs[0].GetShape();
  MxBase::ResizedImageInfo imgInfo;
  imgInfo.widthOriginal = originalImageWidth;
  imgInfo.heightOriginal = originalImageHeight;
  imgInfo.widthResize = YOLOV3_RESIZE;
  imgInfo.heightResize = YOLOV3_RESIZE;
  imgInfo.resizeType = MxBase::RESIZER_STRETCHING;
  std::vector<MxBase::ResizedImageInfo> imageInfoVec = {};
  imageInfoVec.push_back(imgInfo);

  // do postProcess
  postProcessorDptr->Process(tensors, objectInfos, imageInfoVec);

  std::vector<uint32_t> faceId;

  // print result
  std::cout << std::endl;
  std::cout << "Size of objectInfos is " << objectInfos.size() << std::endl;
  std::cout << std::endl;

  for (size_t i = 0; i < objectInfos.size(); i++)
  {
    std::cout << std::endl;
    std::cout << "objectInfo-" << i << " ,Size:" << objectInfos[i].size() << std::endl;
    std::cout << std::endl;

    for (size_t j = 0; j < objectInfos[i].size(); j++)
    {
      if (objectInfos[i][j].className == "face")
      {
        std::cout << std::endl;
        std::cout << "*****objectInfo-" << i << ":" << j << std::endl;
        std::cout << "x0 is " << objectInfos[i][j].x0 << std::endl;
        std::cout << "y0 is " << objectInfos[i][j].y0 << std::endl;
        std::cout << "x1 is " << objectInfos[i][j].x1 << std::endl;
        std::cout << "y1 is " << objectInfos[i][j].y1 << std::endl;
        std::cout << "confidence is " << objectInfos[i][j].confidence << std::endl;
        std::cout << "classId is " << objectInfos[i][j].classId << std::endl;
        std::cout << "className is " << objectInfos[i][j].className << std::endl;
        std::cout << std::endl;

        faceId.push_back(j);
      }
    }
  }

  // get crop rect only for face
  cropRects.resize(faceId.size());
  for (size_t i = 0; i < faceId.size(); i++)
  {
    cropRects[i].x0 = objectInfos[0][faceId[i]].x0;
    cropRects[i].y0 = objectInfos[0][faceId[i]].y0;
    cropRects[i].x1 = objectInfos[0][faceId[i]].x1;
    cropRects[i].y1 = objectInfos[0][faceId[i]].y1;
  }

  return APP_ERR_OK;
};

WheNet::WheNet(const V2Param &v2Param)
{
  deviceId = v2Param.deviceId;
  std::string modelPath = v2Param.modelPath;

  // model init
  modelDptr = std::make_shared<MxBase::Model>(modelPath, deviceId);
  if (modelDptr == nullptr)
  {
    LogError << "modelDptr nullptr";
  }
};

APP_ERROR WheNet::Infer(MxBase::Image &inputImage,
                        std::vector<MxBase::Tensor> &outputs)
{
  APP_ERROR ret;

  // move image to device
  MxBase::Tensor tensorImg = inputImage.ConvertToTensor();
  ret = tensorImg.ToDevice(deviceId);
  if (ret != APP_ERR_OK)
  {
    LogError << "ToDevice failed, ret=" << ret;
    return ret;
  }

  // make infer input
  std::vector<MxBase::Tensor> inputs = {tensorImg};

  // do infer
  outputs = modelDptr->Infer(inputs);

  // move result to host
  for (size_t i = 0; i < outputs.size(); i++)
  {
    outputs[i].ToHost();
  }

  return APP_ERR_OK;
};

void WheNet::Softmax(std::vector<float> src, std::vector<float> &dst)
{
  float maxPosition = *max_element(src.begin(), src.end());
  std::vector<double> a;
  double b = 0;
  for (float &element : src)
  {
    element -= maxPosition;
    a.push_back(exp(element));
    b += exp(element);
  }
  for (double &element : a)
  {
    dst.push_back(element * 1.0 / b);
  }
};

float WheNet::GetAngleFromTensor(MxBase::Tensor &tensor, const int shape_size)
{
  std::cout << std::endl;
  std::cout << "tensor shape size-->[1, " << shape_size << "]" << std::endl;
  auto dataPtr = (float *)tensor.GetData();

  std::vector<float> myangle, angle_predicted_vec;
  for (int i = 0; i < shape_size; i++)
  {
    myangle.push_back(dataPtr[i]);
  }
  Softmax(myangle, angle_predicted_vec);
  float angle_predicted = 0;
  for (int i = 0; i < shape_size; i++)
  {
    angle_predicted += angle_predicted_vec[i] * i;
  }
  // Divide by 2 to turn the angle into the same positive and negative range,so as to avoid being positive numbers
  angle_predicted = angle_predicted * BIN_WIDTH_IN_DEGREES - (shape_size / 2 * BIN_WIDTH_IN_DEGREES);
  return angle_predicted;
};

APP_ERROR WheNet::PostProcess(std::vector<MxBase::Tensor> &outputs, MxBase::Rect cropConfig, std::map<std::string, double> &drawMap)
{
  for (int index = 0; index < (int)outputs.size(); index += TENSORS_PER_OBJECT)
  {
    // Get output shape of model
    int yaw_index = index;
    int pitch_index = index + 1;
    int roll_index = index + 2;
    auto yaw_shape = outputs[yaw_index].GetShape();
    auto pitch_shape = outputs[pitch_index].GetShape();
    auto roll_shape = outputs[roll_index].GetShape();
    // Generate yaw,pitch,roll
    int effective_shape_index = 1;
    // Shape: [1, effective_shape]
    float yaw_predicted = GetAngleFromTensor(outputs[yaw_index], yaw_shape[effective_shape_index]);
    float pitch_predicted = GetAngleFromTensor(outputs[pitch_index], pitch_shape[effective_shape_index]);
    float roll_predicted = GetAngleFromTensor(outputs[roll_index], roll_shape[effective_shape_index]);

    std::cout << std::endl;
    std::cout << "yaw_predicted:" << yaw_predicted << std::endl;
    std::cout << "pitch_predicted:" << pitch_predicted << std::endl;
    std::cout << "roll_predicted:" << roll_predicted << std::endl;
    std::cout << std::endl;

    // save yaw,pitch,roll to file v2_result.txt for test
    std::ofstream output_file("./v2_result.txt", ios::out | ios::app);
    output_file << yaw_predicted << "," << pitch_predicted << "," << roll_predicted << std::endl;
    output_file.close();

    float tdx = (cropConfig.x0 + cropConfig.x1) / 2;
    float tdy = (cropConfig.y0 + cropConfig.y1) / 2;
    int size = 100;
    WheNetDraw(pitch_predicted, yaw_predicted, roll_predicted, tdx, tdy, size, drawMap);
  }
  return APP_ERR_OK;
};

void WheNet::WheNetDraw(float pitch, float yaw, float roll, float tdx, float tdy, int size, std::map<std::string, double> &drawMap)
{
  // Angle to radian: radian = angle * 2pi / circumferential_angle
  int circumferential_angle = 360;
  pitch = pitch * TWO * PI / circumferential_angle;
  yaw = -(yaw * TWO * PI / circumferential_angle);
  roll = roll * TWO * PI / circumferential_angle;
  // X-Axis pointing to right. drawn in red
  double x1 = size * (cos(yaw) * cos(roll)) + tdx;
  double y1 = size * (cos(pitch) * sin(roll) + cos(roll) * sin(pitch) * sin(yaw)) + tdy;
  // Y-Axis | drawn in green
  double x2 = size * (-cos(yaw) * sin(roll)) + tdx;
  double y2 = size * (cos(pitch) * cos(roll) - sin(pitch) * sin(yaw) * sin(roll)) + tdy;
  // Z-Axis (out of the screen) drawn in blue
  double x3 = size * sin(yaw) + tdx;
  double y3 = size * (-cos(yaw) * sin(pitch)) + tdy;
  // map
  drawMap.insert(std::pair<std::string, double>("box_width", tdx));
  drawMap.insert(std::pair<std::string, double>("box_height", tdy));
  drawMap.insert(std::pair<std::string, double>("yaw_x", x1));
  drawMap.insert(std::pair<std::string, double>("yaw_y", y1));
  drawMap.insert(std::pair<std::string, double>("pitch_x", x2));
  drawMap.insert(std::pair<std::string, double>("pitch_y", y2));
  drawMap.insert(std::pair<std::string, double>("roll_x", x3));
  drawMap.insert(std::pair<std::string, double>("roll_y", y3));
};